#include "HMM.h"

namespace ducks
{
    HMM::HMM(int N, int M)
    {
        this->N = N;
        this->M = M;
        this->currT = 0;
        this->validStates = 3;
        this->c = NULL;
        this->alpha = this->beta = this->gamma = NULL;
        this->digamma = NULL;
        this->possibleStates = new int[validStates];
        for(int i = 0; i < validStates; ++ i)
            possibleStates[i] = -1;
        this->acceptState = 0.99;
        this->knownStates = 0;
        this->nextGuessProb = 0;
        
        pi = new float[N];

        A = new float*[N];
        for(int i = 0; i < N; ++ i)
            A[i] = new float[N];

        B = new float*[N];
        for(int i = 0; i < N; ++ i)
            B[i] = new float[M];

        srand(time(NULL));
        float eps;
        float sum = 0;
    
        //Randomly initialize A
        for(int i = 0; i < N; ++ i)
        {
            sum = 0;
            eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)N);
            for(int j = 0; j < N; ++ j)
            {
                A[i][j] = (1.0 / (float)N) + eps;
                if(eps < 0)
                {
                    eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)N);
                }else
                {
                    eps *= -1;
                }
                sum += A[i][j];
            }
            A[i][N - 1] += 1.0 - sum;
        }

        //Randomly initialize B
        for(int i = 0; i < N; ++ i)
        {
            eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)M);
            sum = 0;
            for(int j = 0; j < M; ++ j)
            {
                B[i][j] = (1.0 / (float)M) + eps;
                if(eps < 0)
                {
                    eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)M);
                }else
                {
                    eps *= -1;
                }
                sum += B[i][j];
            }
            B[i][M - 1] += 1.0 - sum;
        }

        //Randomly initialize pi
        eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)N);
        sum = 0;
        for(int i = 0; i < N; ++ i)
        {
            pi[i] = (1.0 / (float)N) + eps;
            if(eps < 0)
            {
                eps = ((float)rand() / (float)RAND_MAX) * (1.0 / (float)N);
            }else
            {
                eps *= -1;
            }
            sum += pi[i];
        }
        pi[N - 1] += 1.0 - sum;
    }

    void HMM::estimate(vector<int> obs, int maxIters)
    {
        float oldLogProb,  newLogProb = FLT_MIN, eps = 0.0000001;
        int T = (int)obs.size(), iters = 0;
        
        //Delete whatever c had and create a new one
        if(this->c != NULL)
            delete[] c;
        this->c = new float[T];

        //Delete any value alpha had
        if(this->alpha != NULL)
        {
            for(int i = 0; i < this->currT; ++ i)
            {
                if(this->alpha[i] != NULL)
                    delete[] alpha[i];
            }
            delete[] alpha;
        }
        this->alpha = new float*[T];
        for(int i = 0; i < T; ++ i)
            this->alpha[i] = new float[this->N];

        //Reallocating the space for beta
        if(this->beta != NULL)
        {
            for(int i = 0; i < this->currT; ++ i)
            {
                if(this->beta[i] != NULL)
                    delete[] (this->beta)[i];
            }
            delete[] this->beta;
        }
        this->beta = new float*[T];
        for(int i = 0; i < T; ++ i)
            (this->beta)[i] = new float[this->N];

        //Reallocating the space for gamma
        if(this->gamma != NULL)
        {
            for(int i = 0; i < (this->currT - 1); ++ i)
            {
                if(this->gamma[i] != NULL)
                    delete[] (this->gamma)[i];
            }
            delete[] this->gamma;
        }
        this->gamma = new float*[T - 1];
        for(int i = 0; i < T - 1; ++ i)
            (this->gamma)[i] = new float[this->N];

        //Reallocating the space for digamma
        if(this->digamma != NULL)
        {
            for(int i = 0; i < (this->currT - 1); ++ i)
            {
                if(this->digamma[i] != NULL)
                {
                    for(int j = 0; j < this->N; ++ j)
                    {
                        if(this->digamma[i][j] != NULL)
                            delete[] (this->digamma)[i][j];
                    }
                    delete[] (this->digamma)[i];
                }
            }
            delete[] this->digamma;
        }
        this->digamma = new float**[T - 1];
        for(int i = 0; i < T - 1; ++ i)
        {
            (this->digamma)[i] = new float*[this->N];
            for(int j = 0; j < this->N; ++ j)
                (this->digamma)[i][j] = new float[this->N];
        }

        this->currT = T;

        do
        {
            oldLogProb = newLogProb;
            //Perform the alpha-pass
            this->alphaPass(obs, this->c, this->alpha);

            //Perform the beta-pass
            //Compute the first betas
            for(int i = 0; i < this->N; ++ i)
                (this->beta)[T - 1][i] = c[T - 1];

            //Compute the rest of the betas
            for(int t = T - 2; t >= 0; -- t)
            {
                for(int i = 0; i < this->N; ++ i)
                {
                    (this->beta)[t][i] = 0;
                    for(int j = 0; j < this->N; ++ j)
                        (this->beta)[t][i] += (this->A)[i][j] * 
                                              (this->B)[j][obs[t + 1]] *
                                              (this->beta)[t + 1][j];

                    //Scale beta with the same scale factor as alpha 
                    (this->beta)[t][i] *= c[t];
                }
            }
            //End of the beta-pass
            
            //Compute gammas and di-gammas
            for(int t = 0; t < T - 1; ++ t)
            {
                float denom = 0;
                for(int i = 0; i < this->N; ++ i)
                {
                    for(int j = 0; j < this->N; ++ j)
                    {
                        denom += (this->alpha)[t][i] *
                                 (this->A)[i][j] *
                                 (this->B)[j][obs[t + 1]] *
                                 (this->beta)[t + 1][j];
                    }
                }
                float maxGamma = 0;
                float likeState = 0;
                for(int i = 0; i < this->N; ++ i)
                {
                    (this->gamma)[t][i] = 0;
                    for(int j = 0; j < this->N; ++ j)
                    {
                        (this->digamma)[t][i][j] = (
                                                    (this->alpha)[t][i] *
                                                    (this->A)[i][j] *
                                                    (this->B)[j][obs[t + 1]] *
                                                    (this->beta)[t + 1][j]
                                                   ) / denom;
                        (this->gamma)[t][i] += (this->digamma)[t][i][j];
                    }
/*                    if(this->knownStates < 3)
                    {
                        if((this->gamma)[t][i] > maxGamma)
                        {
                            maxGamma = (this->gamma)[t][i];
                            likeState = i;
                        }
                    }*/
                }
/*                if(this->knownStates < validStates)
                {
                    bool hasBeen = false;
                    for(int i = 0; i < this->knownStates; ++ i)
                        if(this->possibleStates[i] == likeState)
                            hasBeen = true;
                    if(!hasBeen && (maxGamma > this->acceptState))
                    {
                        possibleStates[knownStates ++] = likeState;
                    }
                    if(knownStates == 3)
                    {
                        this->redoMatrix();
                    }
                }*/
            }

            //Re-estimate the current model
            //Pi will not be reestimated, since there is particular information
            //that will allow better results. Initially it will have a random 
            //distribution, but when the 3 states are defined, the distribution
            //will be fixed by having an equal probability of starting on any
            //of the valid states
/*            for(int i = 0; i < this->N; ++ i)
            {
                (this->pi)[i] = (this->gamma)[0][i];
                cout << pi[i] << " ";
            }
            cout << endl;*/

            //A
            for(int i = 0; i < this->N; ++ i)
                for(int j = 0; j < this->N; ++ j)
                {
                    float numer = 0;
                    float denom = 0;
                    
                    for(int t = 0; t < T - 1; ++ t)
                    {
                        numer += (this->digamma)[t][i][j];
                        denom += (this->gamma)[t][i];
                    }
                    (this->A)[i][j] = numer / denom;
                }
            //B
            for(int i = 0; i < this->N; ++ i)
                for(int j = 0; j < this->M; ++ j)
                {
                    float numer = 0;
                    float denom = 0;
 
                    for(int t = 0; t < T - 1; ++ t)
                    {
                        if(obs[t] == j)
                            numer += (this->gamma)[t][i];
                        denom += (this->gamma)[t][i];
                    }
                    (this->B)[i][j] = numer / denom;
                }
            //Compute log(P(obs|this))
            newLogProb = 0;
            for(int i = 0; i < T; ++ i)
            {
                newLogProb += log(c[i]);
            }

            newLogProb *= -1;
            ++ iters;
        }while((fabs(newLogProb - oldLogProb) > eps) &&
                (iters < maxIters));
    }

    int HMM::calculateNextObservation()
    {
        //First, I estimate which is the most likely state at the moment
        int currState = 0;
        float maxGamma = 0;

        float denom = 0;
        for(int j = 0; j < this-> N; ++ j)
        {
            denom += (this->alpha)[this->currT - 1][j];
        }

        for(int i = 0; i < this->N; ++ i)
        {
            float res = (this->alpha)[this->currT - 1][i] * 
                        (this->beta)[this->currT - 1][i] / 
                        denom;
            if(res > maxGamma)
            {
                maxGamma = res;
                currState = i;
            }
        }

        int nextObs = -1;
        float prob = 0.0;

        for(int t = 0; t < this->M; ++ t)
            if((this->B)[currState][t] > prob)
            {
                prob = (this->B)[currState][t];
                nextObs = t;
            }

        this->nextGuessProb = prob;
        return nextObs;
    }

    void HMM::redoMatrix()
    {
        bool found[this->N];
        float newPi = 1 / (float)(this->N - 1);

        memset(found, 0, (this->N) * sizeof(bool));

        int impState;
        //Determine the impossible state
        for(int i = 0; i < this->validStates; ++ i)
            found[this->possibleStates[i]] = true;

        for(int i = 0; i < this->N; ++ i)
            if(!found[i])
            {
                impState = i;
                break;
            }

        for(int i = 0; i < this->N; ++ i)
        {
            float oldP = this->A[i][impState];
            oldP = oldP / (this->N - 1);
            //Redistribute evenly the probability left by this state to
            //preserve the row stocasthic property
            if(i != impState)
                for(int j = 0; j < this->N; ++ j)
                {
                    if(j == impState)
                        continue;
                    this->A[i][j] += oldP;
                }
            this->A[i][impState] = 0;
            this->A[impState][i] = 0;
            this->pi[i] = newPi;
        }
        this->pi[impState] = 0;
    }

    void HMM::alphaPass(vector<int> obs, float *c, float **ret)
    {
        int T = obs.size();

        c[0] = 0;
        for(int i = 0; i < this->N; ++ i)
        {
            ret[0][i] = (this->pi)[i] * (this->B)[i][obs[0]];
            c[0] += ret[0][i];                
        }

        //Scale the first alphas
        c[0] = 1 / c[0];
        for(int i = 0; i < this->N; ++ i)
        {
            ret[0][i] *= c[0];
        }
        //Compute the rest of the alphas
        for(int t = 1; t < T; ++ t)
        {
            c[t] = 0;
            for(int i = 0; i < this->N; ++ i)
            {
                ret[t][i] = 0;
                for(int j = 0; j < this->N; ++ j)
                {
                    ret[t][i] += ret[t - 1][j] *
                                 (this->A)[j][i];
                }

                ret[t][i] *= (this->B)[i][obs[t]];
                c[t] += ret[t][i];
            }


            //Scale each alpha
            c[t] = 1 / c[t];
            for(int i = 0; i < this->N; ++ i)
            {
                ret[t][i] *= c[t];
 
            }
        }
    }

    float HMM::probabilityOf(vector<int> obs)
    {
        float dummyC[obs.size()];
        float ret = 0;
        int size = (this->currT < obs.size() ? this->currT : obs.size());
        float **alph;
        vector<int> toTest;

        //If the observation is larger than the one that built this model,
        //trimm it to avoid a response of 0 due to the overfitting. Else it
        //will be completely added
        for(int i = 0; i < size; ++ i)
            toTest.push_back(obs[i]);

        alph = new float*[size];
        for(int i = 0; i < size; ++ i)
            alph[i] = new float[this->N];
        this->alphaPass(toTest, dummyC, alph);

        for(int i = 0; i < this->N; ++ i)
        {
            ret += (alph[size - 1][i] / dummyC[size - 1]);
        }

        //Set free the alphas
        for(int i = 0; i < size; ++ i)
            delete[] alph[i];
        delete[] alph;

        return ret;
    }
    HMM::~HMM()
    {
        //First, delete all the matrices that are built by the constructor
        //Delete A
        for(int i = 0; i < this->N; ++ i)
        {
            delete[] this->A[i];
        }
        delete[] this->A;

        //Delete B
        for(int i = 0; i < this->N; ++ i)
        {
            delete[] this->B[i];
        }
        delete[] this->B;

        //Delete pi
        delete[] this->pi;

        //Delete possibleStates
        delete[] this->possibleStates;

        //Now, delete the matrices that my or may not be initialized
        if(this->c != NULL)
            delete[] this->c;

        //Delete alpha
        if(this->alpha != NULL)
        {
            for(int i = 0; i < this->currT; ++ i)
            {
                if(this->alpha[i] != NULL)
                    delete[] alpha[i];
            }
            delete[] alpha;
        }

        //Delete beta
        if(this->beta != NULL)
        {
            for(int i = 0; i < this->currT; ++ i)
            {
                if(this->beta[i] != NULL)
                    delete[] (this->beta)[i];
            }
            delete[] this->beta;
        }

        //Delete gamma
        if(this->gamma != NULL)
        {
            for(int i = 0; i < (this->currT - 1); ++ i)
            {
                if(this->gamma[i] != NULL)
                    delete[] (this->gamma)[i];
            }
            delete[] this->gamma;
        }

        //Delete digamma
        if(this->digamma != NULL)
        {
            for(int i = 0; i < (this->currT - 1); ++ i)
            {
                if(this->digamma[i] != NULL)
                {
                    for(int j = 0; j < this->N; ++ j)
                    {
                        if(this->digamma[i][j] != NULL)
                            delete[] (this->digamma)[i][j];
                    }
                    delete[] (this->digamma)[i];
                }
            }
            delete[] this->digamma;
        }
    } 
}
